# Author : ZZH
# Date : 2025/5/21
import argparse
import platform
import shutil

import psutil
import yaml
from ultralytics import YOLO
from datetime import datetime
import os
import logging
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'


def parse_args():
    """
    解析命令行参数，配置训练所需要的超参数
    :return: argparse.Namespace 包含所有命令行的参数对象
    """
    parser = argparse.ArgumentParser(description="YOLOv8增强训练脚本")
    # 配置文件参数
    parser.add_argument("--model", type=str, default='yolov8n', help="模型名称（如yolov8n）")
    parser.add_argument("--data", type=str, default='../configs/data.yaml', help="数据集配置文件名（自动查找config/）")

    # 训练参数
    parser.add_argument("--epochs", type=int, default=2, help="训练轮次")
    parser.add_argument("--batch", type=int, default=16, help="批次大小")
    parser.add_argument("--imgsz", type=int, default=640, help="图像尺寸")
    parser.add_argument("--device", type=str, default="cpu", help="训练设备")
    parser.add_argument("--lr0", type=float, default=0.01, help="初始学习率")
    parser.add_argument("--lr", type=float, default=None, help="最终学习率")
    parser.add_argument("--run-name", type=str, default="train", help="运行名称")
    return parser.parse_args()


def setup_logging(args, data_info):
    """初始化符合项目要求的日志系统"""
    # 生成标准化日志路径
    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
    log_dir = os.path.join("..", "logging", "detect")
    os.makedirs(log_dir, exist_ok=True)
    log_path = os.path.join(log_dir, f"train_{timestamp}_{args.model}.log")

    # 初始化日志记录器
    logger = logging.getLogger("YOLO_TRAIN")
    logger.setLevel(logging.INFO)

    # 配置文件处理器
    file_handler = logging.FileHandler(log_path, encoding='utf-8-sig')
    formatter = logging.Formatter(
        '[%(asctime)s] %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S'
    )
    file_handler.setFormatter(formatter)
    logger.addHandler(file_handler)

    # 记录必要信息
    logger.info("[训练参数] epochs=%d, batch=%d, lr0=%.4f, imgsz=%d",
                args.epochs, args.batch, args.lr0, args.imgsz)
    logger.info("[数据集信息] 类别数: %d (%s)",
                data_info['nc'], '、'.join(data_info['names']))
    logger.info("[训练集路径] %s", data_info['train'])

    # 硬件信息记录
    logger.info("[硬件信息] OS: %s %s",
                platform.system(), platform.release())
    logger.info("[硬件信息] CPU: %s (Cores: %d)",
                platform.processor(), psutil.cpu_count(logical=False))

    # GPU信息（需要安装torch）
    try:
        import torch
        if torch.cuda.is_available():
            logger.info("[硬件信息] GPU: %s (VRAM: %.1fGB)",
                        torch.cuda.get_device_name(0),
                        torch.cuda.get_device_properties(0).total_memory / 1e9)
            logger.info("[硬件信息] CUDA: %s", torch.version.cuda)
        else:
            logger.info("[硬件信息] GPU: 未检测到可用显卡")
    except ImportError:
        logger.warning("[硬件信息] 无法获取GPU信息，请安装PyTorch")

    return log_path


def find_config(filename):
    """自动查找配置文件"""
    config_path = os.path.join("..", "configs", filename)
    if not os.path.exists(config_path):
        raise FileNotFoundError(f"配置文件 {filename} 未找到于 {os.path.dirname(config_path)}")
    return config_path


def handle_pretrained_model(model_name):
    """处理预训练模型逻辑"""
    pretrained_dir = os.path.join("..", "models", "pretrained")
    model_path = os.path.join(pretrained_dir, f"{model_name}.pt")

    # 模型不存在时自动下载
    if not os.path.exists(model_path):
        print(f"预训练模型 {model_name}.pt 未找到，开始下载...")
        model = YOLO(f"{model_name}.pt")
        os.makedirs(pretrained_dir, exist_ok=True)
        shutil.move(f"{model_name}.pt", model_path)
    return model_path


def save_checkpoints(model_name, src_path, is_best=True):
    """保存模型到checkpoints目录"""
    checkpoints_dir = os.path.join("..", "models", "checkpoints")
    os.makedirs(checkpoints_dir, exist_ok=True)

    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
    suffix = "best" if is_best else "last"
    dest_file = f"{timestamp}-{model_name}-{suffix}.pt"
    shutil.copy(src_path, os.path.join(checkpoints_dir, dest_file))


def main():
    args = parse_args()
    # 加载数据集配置信息
    with open(find_config(args.data), 'r', encoding='utf-8') as f:
        data_config = yaml.safe_load(f)

    # 初始化日志系统（新增）
    log_path = setup_logging(args, {
        'nc': data_config['nc'],
        'names': data_config['names'],
        'train': data_config['train']
    })
    print(f"训练日志已保存至：{os.path.abspath(log_path)}")

    # 自动查找配置
    data_config = find_config(args.data)
    model_path = handle_pretrained_model(args.model)

    # 初始化模型
    model = YOLO(model_path)

    # 配置训练参数
    train_args = {
        "data": data_config,
        "epochs": args.epochs,
        "batch": args.batch,
        "imgsz": args.imgsz,
        "device": args.device,
        "lr0": args.lr0,
        "project": os.path.join("..", "runs", "detect"),
        "name": args.run_name or datetime.now().strftime("%Y%m%d%H%M"),
        "exist_ok": True
    }
    if args.lr:
        train_args["lr"] = args.lr

    # 开始训练
    results = model.train(**train_args)

    # 处理输出模型
    weights_dir = os.path.join(train_args["project"], train_args["name"], "weights")
    save_checkpoints(args.model, os.path.join(weights_dir, "best.pt"))
    save_checkpoints(args.model, os.path.join(weights_dir, "last.pt"), is_best=False)


if __name__ == "__main__":
    main()
